#include "sut.h"

#include <queue>
#include <mutex>
#include <atomic>
#include <thread>
#include <cstring>
#include <functional>
#include <condition_variable>

#include <unistd.h>
#include "context.h"
#include <arpa/inet.h>
#include <sys/socket.h>

constexpr std::size_t kStackSizeDefault = 1024 * 4;
constexpr std::size_t kMaxMessageLength = 1024;

using ContextType = SimpleUserStack;

class SimpleUserThreadScheduler;

class SimpleUserThread
{
    ContextType uc;
    int fdRead, fdWrite;
    bool finished = false;
    std::function<void()> task;
    SimpleUserThreadScheduler* sched;
    static SimpleUserThread *sutRunning;
    
    SimpleUserThread(SimpleUserThreadScheduler* s,
                     std::function<void()> t) : sched(s), task(t)
    {
        SetRWDescriptor(fileno(stdin), fileno(stdout));
        memset(&uc, 0, sizeof(JAMScriptUserContext));
        uc.uc_stack.ss_sp = new char[kStackSizeDefault];
        uc.uc_stack.ss_size = kStackSizeDefault;
        CreateContext(&(uc), CommonThreadFunction);
    }
    
    static void CommonThreadFunction()
    {
        auto& t = SimpleUserThread::GetCurrentTask();
        t.task();
        t.OnFinish();
    }
    
public:
    
    template <typename Fn, typename ...Args>
    static bool Create(SimpleUserThreadScheduler *s, Fn&& fn, Args&&... args);
    
    static SimpleUserThread& GetCurrentTask()
    {
        return *sutRunning;
    }

    int GetWriteDescriptor()
    {
        return fdWrite;
    }
    
    int GetReadDescriptor()
    {
        return fdRead;
    }
    
    void SetRWDescriptor(int fdr, int fdw)
    {
        fdRead = fdr;
        fdWrite = fdw;
    }
    
    bool Finished()
    {
        return finished;
    }
    
    void OnFinish()
    {
        finished = true;
        SwapOut();
    }
    
    auto& GetScheduler()
    {
        return *sched;
    }
    
    ~SimpleUserThread()
    {
        delete[] reinterpret_cast<char*>(uc.uc_stack.ss_sp);
    }
    
    void SwapTo();
    void SwapOut();

};

class SimpleUserThreadScheduler
{
    ContextType sc;
    std::atomic_int count;
    std::atomic_bool cont;
    std::mutex mtxIO, mtxTask;
    std::deque<SimpleUserThread *> taskReadyQueue;
    std::deque<std::function<void()>> taskIOQueue;
    
    static SimpleUserThreadScheduler *singleScheduler;
    std::thread ioThread, execThread;
    
    SimpleUserThreadScheduler() : cont(true)
    {
        
    }
    
public:
    
    static SimpleUserThreadScheduler& GetInstance()
    {
        if (singleScheduler == nullptr)
        {
            singleScheduler = new SimpleUserThreadScheduler();
        }
        return *singleScheduler;
    }
    
    static void Cleanup()
    {
        if (singleScheduler != nullptr)
        {
            delete singleScheduler;
            singleScheduler = nullptr;
        }
    }
    
    void OnCreateTask()
    {
        count++;
    }
    
    void Start()
    {
        execThread = std::thread([this] {
            this->ComputeExec();
        });
        ioThread = std::thread([this] {
            this->InputOutputExec();
        });
    }
    
    void ShutDown()
    {
        cont = false;
        execThread.join();
        ioThread.join();
    }
    
    ContextType& GetSchedulerContext()
    {
        return sc;
    }
    
    SimpleUserThread *GetNextTask()
    {
        std::scoped_lock lk(mtxTask);
        if (!taskReadyQueue.empty())
        {
            auto nextTask = taskReadyQueue.front();
            taskReadyQueue.pop_front();
            return nextTask;
        }
        return nullptr;
    }
    
    void ComputeExec()
    {
        while (cont || count > 0)
        {
            auto *nextTask = GetNextTask();
            if (nextTask != nullptr)
            {
                nextTask->SwapTo();
                if (nextTask->Finished())
                {
                    count--;
                    delete nextTask;
                }
            }
            else
            {
                std::this_thread::sleep_for(std::chrono::microseconds(10));
            }
        }
    }

    void InputOutputExec()
    {
        while (cont || count > 0)
        {
            std::unique_lock lk(mtxIO);
            if (!taskIOQueue.empty())
            {
                auto& ioFunctor = taskIOQueue.front();
                lk.unlock();
                ioFunctor();
                lk.lock();
                taskIOQueue.pop_front();
            }
            else
            {
                lk.unlock();
                std::this_thread::sleep_for(std::chrono::microseconds(10));
            }
        }
    }
    
    void PostComputeTask(SimpleUserThread *t)
    {
        std::unique_lock lk(mtxTask);
        taskReadyQueue.push_back(t);
    }

    void PostIOTask(std::function<void()> tio)
    {
        std::unique_lock lk(mtxIO);
        taskIOQueue.push_back(tio);
    }
    
};

SimpleUserThreadScheduler *SimpleUserThreadScheduler::singleScheduler = nullptr;
SimpleUserThread *SimpleUserThread::sutRunning = nullptr;

void SimpleUserThread::SwapTo()
{
    sutRunning = this;
    SwapToContext(&(sched->GetSchedulerContext()), &uc, nullptr);
}

void SimpleUserThread::SwapOut()
{
    sutRunning = nullptr;
    SwapToContext(&uc, &(sched->GetSchedulerContext()), nullptr);
}

template <typename Fn, typename ...Args>
bool SimpleUserThread::Create(SimpleUserThreadScheduler *s,
                              Fn&& fn, Args&&... args)
{
    auto pNewTask = new SimpleUserThread(s,
                                [fn = std::forward<Fn>(fn),
                                 ...args = std::forward<Args>(args)] ()
                                mutable {
        fn(std::forward<Args>(args)...);
    });
    if (pNewTask != nullptr)
    {
        SimpleUserThreadScheduler::GetInstance().OnCreateTask();
        SimpleUserThreadScheduler::GetInstance().PostComputeTask(pNewTask);
        return true;
    }
    return false;
}

template <typename Fn, typename ...Args>
void IOExecutor(Fn&& fn, Args&& ...args)
{
    auto& taskRunningIO = SimpleUserThread::GetCurrentTask();
    taskRunningIO
    .GetScheduler()
    .PostIOTask([fn = std::forward<Fn>(fn),
                 ...args = std::forward<Args>(args),
                 &taskRunningIO] () mutable {
        fn(std::forward<Args>(args)...);
        taskRunningIO.GetScheduler().PostComputeTask(&taskRunningIO);
    });
    taskRunningIO.SwapOut();
}

void sut_init()
{
    SimpleUserThreadScheduler::GetInstance().Start();
}

bool sut_create(sut_task_f fn)
{
    return SimpleUserThread::Create(&(SimpleUserThreadScheduler::GetInstance()), fn);
}

void sut_yield()
{
    auto& toYield = SimpleUserThread::GetCurrentTask();
    SimpleUserThreadScheduler::GetInstance().PostComputeTask(&toYield);
    return toYield.SwapOut();
}

void sut_exit()
{
    return SimpleUserThread::GetCurrentTask().OnFinish();
}

void sut_open(char *dest, int port)
{
    auto& currTask = SimpleUserThread::GetCurrentTask();
    return IOExecutor([&currTask, dest, port] () mutable {
        struct sockaddr_in server;
        int fdSocket = socket(AF_INET, SOCK_STREAM, 0);
        if (fdSocket == -1)
        {
            perror("[SUT-ERROR]: Socket Creation");
            return;
        }
        server.sin_addr.s_addr = inet_addr(dest);
        server.sin_family = AF_INET;
        server.sin_port = htons(port);
        if (connect(fdSocket, (struct sockaddr *)&server, sizeof(server)) < 0)
        {
            perror("[SUT-ERROR]: Connection");
            return;
        }
        currTask.SetRWDescriptor(fdSocket, fdSocket);
    });
}

void sut_write(char *buf, int size)
{
    return IOExecutor(write, SimpleUserThread::GetCurrentTask().GetWriteDescriptor(), buf, size);
}

void sut_close()
{
    close(SimpleUserThread::GetCurrentTask().GetReadDescriptor());
    SimpleUserThread::GetCurrentTask().SetRWDescriptor(fileno(stdin), fileno(stdout));
}

char *sut_read()
{
    char *res = reinterpret_cast<char*>(calloc(kMaxMessageLength, sizeof(char)));
    auto& currTask = SimpleUserThread::GetCurrentTask();
    IOExecutor(read, currTask.GetReadDescriptor(), res, kMaxMessageLength);
    return res;
}

void sut_shutdown()
{
    SimpleUserThreadScheduler::GetInstance().ShutDown();
    SimpleUserThreadScheduler::Cleanup();
}
